{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image-to-Image Task Guide"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Image-to-Image タスクは、アプリケーションが画像を受信し、別の画像を出力するタスクです。これには、画像強化 (超解像度、低光量強化、ディレインなど)、画像修復などを含むさまざまなサブタスクがあります。\n",
    "\n",
    "このガイドでは、次の方法を説明します。\n",
    "- 超解像度タスクに画像間のパイプラインを使用します。\n",
    "- パイプラインを使用せずに、同じタスクに対してイメージ間モデルを実行します。\n",
    "\n",
    "このガイドがリリースされた時点では、`image-to-image`パイプラインは超解像度タスクのみをサポートしていることに注意してください。\n",
    "\n",
    "必要なライブラリをインストールすることから始めましょう。\n",
    "\n",
    "```bash\n",
    "pip install transformers\n",
    "```\n",
    "\n",
    "[Swin2SR モデル](https://huggingface.co/caidas/swin2SR-lightweight-x2-64) を使用してパイプラインを初期化できるようになりました。次に、イメージを使用してパイプラインを呼び出すことで、パイプラインを推論できます。現時点では、[Swin2SR モデル](https://huggingface.co/models?sort=trending&search=swin2sr) のみがこのパイプラインでサポートされています。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "pipe = pipeline(task=\"image-to-image\", model=\"caidas/swin2SR-lightweight-x2-64\", device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "では、画像を読み込みましょう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import requests\n",
    "\n",
    "url = \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/cat.jpg\"\n",
    "image = Image.open(requests.get(url, stream=True).raw)\n",
    "\n",
    "print(image.size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```bash\n",
    "# (532, 432)\n",
    "```\n",
    "<div class=\"flex justify-center\">\n",
    "     <img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/cat.jpg\" alt=\"Photo of a cat\"/>\n",
    "</div>\n",
    "\n",
    "\n",
    "これで、パイプラインを使用して推論を実行できるようになりました。猫の画像の拡大バージョンを取得します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upscaled = pipe(image)\n",
    "print(upscaled.size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```bash\n",
    "# (1072, 880)\n",
    "```\n",
    "\n",
    "パイプラインを使用せずに自分で推論を実行したい場合は、トランスフォーマーの `Swin2SRForImageSuperResolution` クラスと `Swin2SRImageProcessor` クラスを使用できます。これには同じモデルのチェックポイントを使用します。モデルとプロセッサを初期化しましょう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Swin2SRForImageSuperResolution, Swin2SRImageProcessor \n",
    "\n",
    "model = Swin2SRForImageSuperResolution.from_pretrained(\"caidas/swin2SR-lightweight-x2-64\").to(device)\n",
    "processor = Swin2SRImageProcessor(\"caidas/swin2SR-lightweight-x2-64\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`pipeline`」は、自分で行う必要がある前処理と後処理のステップを抽象化するので、画像を前処理しましょう。画像をプロセッサに渡してから、ピクセル値を GPU に移動します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pixel_values = processor(image, return_tensors=\"pt\").pixel_values\n",
    "print(pixel_values.shape)\n",
    "\n",
    "pixel_values = pixel_values.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "これで、ピクセル値をモデルに渡すことで画像を推測できるようになりました。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "with torch.no_grad():\n",
    "  outputs = model(pixel_values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "出力は、以下のような `ImageSuperResolutionOutput` タイプのオブジェクトです 👇\n",
    "\n",
    "```\n",
    "(loss=None, reconstruction=tensor([[[[0.8270, 0.8269, 0.8275,  ..., 0.7463, 0.7446, 0.7453],\n",
    "          [0.8287, 0.8278, 0.8283,  ..., 0.7451, 0.7448, 0.7457],\n",
    "          [0.8280, 0.8273, 0.8269,  ..., 0.7447, 0.7446, 0.7452],\n",
    "          ...,\n",
    "          [0.5923, 0.5933, 0.5924,  ..., 0.0697, 0.0695, 0.0706],\n",
    "          [0.5926, 0.5932, 0.5926,  ..., 0.0673, 0.0687, 0.0705],\n",
    "          [0.5927, 0.5914, 0.5922,  ..., 0.0664, 0.0694, 0.0718]]]],\n",
    "       device='cuda:0'), hidden_states=None, attentions=None)\n",
    "```\n",
    "\n",
    "`reconstruction`を取得し、それを視覚化するために後処理する必要があります。どのように見えるか見てみましょう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs.reconstruction.data.shape\n",
    "# torch.Size([1, 3, 880, 1072])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "出力を圧縮して軸 0 を削除し、値をクリップしてから、それを numpy float に変換する必要があります。次に、軸を [1072, 880] の形状になるように配置し、最後に出力を範囲 [0, 255] に戻します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# squeeze, take to CPU and clip the values\n",
    "output = outputs.reconstruction.data.squeeze().cpu().clamp_(0, 1).numpy()\n",
    "# rearrange the axes\n",
    "output = np.moveaxis(output, source=0, destination=-1)\n",
    "# bring values back to pixel values range\n",
    "output = (output * 255.0).round().astype(np.uint8)\n",
    "Image.fromarray(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"flex justify-center\">\n",
    "     <img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/cat_upscaled.png\" alt=\"Upscaled photo of a cat\"/>\n",
    "</div>"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
